import os
import time

import cv2
import piexif
import numpy as np
from PIL import Image
from flask_moudle.service.predict import yolov8_seg


def analysis(img, size=None):
    """

    :param input_file: Image
    :param size:
    :return: numpy [[[]]]]
    """
    if size is None:
        size = [512, 512]
    # img = Image.open(input_file)
    new_img = Image.new('RGB', (img.size[0], img.size[1]))
    width, height = img.size
    w = width // size[0] + 1
    h = height // size[1] + 1
    print(f'正在识别图像...')
    for i in range(0, w):
        for j in range(0, h):
            box = (i * size[0], j * size[1], (i + 1) * size[0], (j + 1) * size[1])
            if i + 1 == w:
                box = (width - size[0], j * size[1], width, (j + 1) * size[1])
            if j + 1 == h:
                box = (i * size[0], height - size[1], (i + 1) * size[0], height)
            if i + 1 == w and j + 1 == h:
                box = (width - size[0], height - size[1], width, height)
            cropped_img = img.crop(box)
            # split_name = os.path.join(output_folder,
            #                           f"{filename[:-4]}_{j + 1:02d}_{i + 1:02d}" + os.path.splitext(filename)[1])
            r_image = yolov8_seg(cropped_img)[0]
            new_img.paste(r_image, box)
    print(f'识别完成')
    return new_img

def remove_width_black(img, start_w):
    # 获取图像的高度和宽度
    height, width, _ = img.shape
    row_mask = np.all(img[:, start_w:, :] == [0, 0, 0], axis=2)  # 创建掩码，找到所有像素值为[0,0,0]的行
    # if np.any(row_mask):  # 如果存在这样的行
    #     image_array = np.delete(image_array, row_mask, axis=1)  # 删除这些行
    row_mask_1 = np.all(row_mask[:, :] == True, axis=0)
    row_mask_1 = ~row_mask_1
    right = img[:, start_w:, :][:, row_mask_1, :]
    new_jpg = np.zeros((height, right.shape[1] + start_w, 3), dtype=np.uint8)
    new_jpg[:, :start_w, :] = img[:, :start_w, :]
    new_jpg[:, start_w:, :] = right
    return new_jpg


def remove_height_black(img, start_h):
    # 获取图像的高度和宽度
    height, width, _ = img.shape
    row_mask = np.all(img[start_h:, :, :] == [0, 0, 0], axis=2)  # 创建掩码，找到所有像素值为[0,0,0]的行
    row_mask = np.all(row_mask[:, :] == True, axis=1)
    row_mask = ~row_mask
    bottom = img[start_h:, :, :][:, row_mask, :]
    new_jpg = np.zeros((height + start_h, bottom.shape[1], 3), dtype=np.uint8)
    new_jpg[start_h:, :, :] = img[start_h:, :, :]
    new_jpg[start_h:, :, :] = bottom
    return new_jpg


def rotate_bound_white_bg(image, angle):
    (h, w) = image.shape[:2]
    (cX, cY) = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D((cX, cY), -angle, 1.0)
    cos = np.abs(M[0, 0])
    sin = np.abs(M[0, 1])
    nW = int((h * sin) + (w * cos))
    nH = int((h * cos) + (w * sin))
    M[0, 2] += (nW / 2) - cX
    M[1, 2] += (nH / 2) - cY
    # borderValue 缺失背景填充色彩，此处为白色，可自定义
    return cv2.warpAffine(image, M, (nW, nH), borderValue=(0, 0, 0))
    # borderValue 缺省，默认是黑色（0, 0 , 0）


def paste_img(bottom, top):
    mask = bottom[:, :, :] == [0, 0, 0]
    # mask = top[:, :, :] != [0, 0, 0]
    bottom[mask] = top[mask]
    return bottom


def compose_same_img(left_img, right_img):
    hessian = 100
    surf = cv2.xfeatures2d.SIFT_create(hessian)  # 将Hessian Threshold设置为400,阈值越大能检测的特征就越少
    kp1, des1 = surf.detectAndCompute(left_img, None)  # 查找关键点和描述符
    kp2, des2 = surf.detectAndCompute(right_img, None)

    FLANN_INDEX_KDTREE = 0  # 建立FLANN匹配器的参数
    indexParams = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)  # 配置索引，密度树的数量为5
    searchParams = dict(checks=50)  # 指定递归次数
    # FlannBasedMatcher：是目前最快的特征匹配算法（最近邻搜索）
    flann = cv2.FlannBasedMatcher(indexParams, searchParams)  # 建立匹配器
    matches = flann.knnMatch(des2, des1, k=2)  # 得出匹配的关键点

    good = []
    # 提取优秀的特征点
    start_time_filter = time.time()
    for m, n in matches:
        if m.distance < 0.85 * n.distance:  # 如果第一个邻近距离比第二个邻近距离的0.7倍小，则保留
            good.append(m)
    end_time_filter = time.time()
    print("滤波花费的时间:", end_time_filter - start_time_filter, "秒")
    src_pts = np.array([kp2[m.queryIdx].pt for m in good])  # 查询图像的特征描述子索引
    dst_pts = np.array([kp1[m.trainIdx].pt for m in good])  # 训练(模板)图像的特征描述子索引
    H = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 4.0)  # 生成变换矩阵
    h1, w1 = right_img.shape[:2]
    shft = np.array([[1.0, 0, w1], [0, 1.0, h1], [0, 0, 1.0]])
    M = np.dot(shft, H[0])  # 获取左边图像到右边图像的投影映射关系
    dst_corners = cv2.warpPerspective(right_img, M, (w1 * 3, h1 * 3))  # 透视变换，新图像可容纳完整的两幅图
    return dst_corners[0:h1, w1:2 * w1, :], remove_width_black(dst_corners[h1:2 * h1, w1:3 * w1, :], w1), dst_corners[
                                                                                                          2 * h1:3 * h1,
                                                                                                          w1:2 * w1, :]


# 左边为大图像，未处理
def compose_img(compose, image):
    if np.all(compose == image):
        return compose
    height, width, _ = image.shape
    top, middle, bottom = compose_same_img(compose[:, -width:, :], image)
    compose_height, compose_width, _ = compose.shape
    add_w = middle.shape[1] - width
    if add_w > 0:
        new_img = np.zeros((compose_height, compose_width + add_w, 3), dtype=np.uint8)
        new_img[:compose_height, :compose_width, :] = compose
        compose = new_img
    compose[:, -middle.shape[1]:, :] = paste_img(compose[:, -middle.shape[1]:, :], middle[:, :, :])
    return compose


def rotate_jpg(file_dir,angle=90):
    for filename in os.listdir(file_dir):
        file=os.path.join(file_dir,filename)
        tags = piexif.load(file)
        image=cv2.imread(file)
        image=rotate_bound_white_bg(image,angle)
        cv2.imwrite(file,image)
        piexif.insert(piexif.dump(tags),file)


